#!/usr/bin/env python
# -*- coding: utf-8 -*-


"""
Created on 2018年8月4日
@author: Irony
@site: https://pyqt.site , https://github.com/PyQt5
@email: 892768447@qq.com
@file: QListView.显示自定义Widget
@description:
"""

try:
    from PyQt5.QtCore import QSize, pyqtSignal
    from PyQt5.QtGui import QStandardItemModel, QStandardItem, QImage, QPixmap
    from PyQt5.QtWidgets import QMessageBox, QFileDialog, QGridLayout, QDialog, QListView, QWidget, QHBoxLayout, QLineEdit, \
        QPushButton, QApplication, QLabel, QSpinBox, QTextBrowser, QInputDialog
except ImportError:
    from PySide2.QtCore import QSize
    from PySide2.QtGui import QStandardItemModel, QStandardItem
    from PySide2.QtWidgets import QListView, QDialog, QWidget, QHBoxLayout, QLineEdit, \
        QPushButton, QApplication

from ui.CameraCollect import Ui_Dialog as CameraCollect

import cv2
import numpy as np
import threading
import time
import os


class RenameFileDialog(QDialog):

    image_dir = ''
    label_dir = ''

    def __init__(self, *args, **kwargs):
        super(RenameFileDialog, self).__init__(*args, **kwargs)
        # self.setupUi(self)
        self.setWindowTitle("批量重命名文件-图片标注处理")
        self.setUi()

    def setUi(self):
        self.pushButton_select_dir_image = QPushButton('选择image文件夹')
        self.pushButton_select_dir_image.clicked.connect(self.selectImageDir)
        self.pushButton_select_dir_label = QPushButton('选择label文件夹')
        self.pushButton_select_dir_label.clicked.connect(self.selectLabelDir)

        self.pushButton_rename = QPushButton('批量重命名')
        self.pushButton_rename.clicked.connect(self.renameFile)
        self.pushButton_rename_map = QPushButton('image->label对应批量重命名')
        self.pushButton_rename_map.clicked.connect(self.renameFileMap)
        self.label_image_dir = QLabel('image路径')
        self.label_label_dir = QLabel('label路径')
        self.textBowser_image = QTextBrowser()
        self.textBowser_label = QTextBrowser()

        self.layout = QGridLayout()
        self.layout.addWidget(self.pushButton_select_dir_image, 0, 0, 1, 1)
        self.layout.addWidget(self.label_image_dir, 1, 0, 1, 1)
        self.layout.addWidget(self.pushButton_rename, 2, 0, 1, 1)
        self.layout.addWidget(self.textBowser_image, 3, 0, 1, 1)

        self.layout.addWidget(self.pushButton_select_dir_label, 0, 1, 1, 1)
        self.layout.addWidget(self.label_label_dir, 1, 1, 1, 1)
        self.layout.addWidget(self.pushButton_rename_map, 2, 1, 1, 1)
        self.layout.addWidget(self.textBowser_label, 3, 1, 1, 1)
        self.setLayout(self.layout)

    def selectImageDir(self):
        res = QFileDialog.getExistingDirectory(self, "选择保存路径", "./")
        if res != '':
            self.label_image_dir.setText(res)
            self.image_dir = res
            files = os.listdir(res)
            for f in files:
                self.textBowser_image.append(f)

    def selectLabelDir(self):
        res = QFileDialog.getExistingDirectory(self, "选择保存路径", "./")
        if res != '':
            self.label_label_dir.setText(res)
            self.label_dir = res
            files = os.listdir(res)
            for f in files:
                self.textBowser_label.append(f)

    def renameFile(self):
        path = self.image_dir
        if not path:
            QMessageBox.warning(self, '错误', '请选择image文件夹', QMessageBox.Ok)
            return
        start_num, okPressed = QInputDialog.getText(self, '输入编号', '起始编号:', QLineEdit.Normal, '')
        if not okPressed or not start_num:
            return
        if not start_num.isdigit():
            return
        start_num = int(start_num)

        files = os.listdir(path)
        for i in range(len(files)):
            f = files[i]
            end = f.split('.')[1]
            number = start_num + i
            new_name = str(number).zfill(12) + '.' + end

            old_path = '%s/%s' % (path, f)
            new_path = '%s/%s' % (path, new_name)
            os.rename(old_path, new_path)
            self.textBowser_image.append('rename: %s -> %s' % (old_path.split('/')[-1], new_path.split('/')[-1]))

    def renameFileMap(self):
        # 重命名，image->label 映射关系
        image_dir = self.image_dir
        if not image_dir:
            QMessageBox.warning(self, '错误', '请选择image文件夹', QMessageBox.Ok)
            return
        label_dir = self.label_dir
        if not label_dir:
            QMessageBox.warning(self, '错误', '请选择label文件夹', QMessageBox.Ok)
            return
        start_num, okPressed = QInputDialog.getText(self, '输入编号', '起始编号:', QLineEdit.Normal, '')
        if not okPressed or not start_num:
            return
        if not start_num.isdigit():
            return
        start_num = int(start_num)

        image_files = os.listdir(image_dir)
        label_files = os.listdir(label_dir)
        for i, f in enumerate(image_files):
            name = f.split('.')[0]
            end = f.split('.')[1]

            old_path = '%s/%s' % (image_dir, f)
            new_path = '%s/%s' % (image_dir, str(start_num + i).zfill(12) + '.' + end)
            os.rename(old_path, new_path)
            self.textBowser_image.append('rename: %s -> %s' % (old_path.split('/')[-1], new_path.split('/')[-1]))

            for label in label_files:
                f_label = name + '.txt'
                if f_label == label:
                    old_path = '%s/%s' % (label_dir, f_label)
                    new_path = '%s/%s' % (label_dir, str(start_num + i).zfill(12) + '.txt')
                    os.rename(old_path, new_path)
                    self.textBowser_label.append('rename: %s -> %s' % (old_path.split('/')[-1], new_path.split('/')[-1]))
                    break

    def createDir(self, path):
        isExists = os.path.exists(path)
        if not isExists:
            os.makedirs(path)
            return True
        else:
            # 如果目录存在则不创建，并提示目录已存在
            return False


if __name__ == '__main__':
    import sys

    app = QApplication(sys.argv)
    w = RenameFileDialog()
    w.show()
    sys.exit(app.exec_())
